/**
 * 
 */
package edu.umd.clip.lm.ngram;

import java.io.Serializable;
import java.util.*;

import edu.umd.clip.lm.util.*;


/**
 * @author Denis Filimonov <den@cs.umd.edu>
 *
 */
public class Counts implements Serializable {
	private static final long serialVersionUID = 1L;

	static class Count implements Serializable {
		private static final long serialVersionUID = 1L;

		Long2IntMap counts;

		public Count(Long2IntMap counts) {
			this.counts = counts;
		}
		
		public Count() {
			this.counts = new Long2IntMap();
		}
		
		public void addCount(int word, int count) {
			counts.addAndGet(word, count);
		}
	}
	
	private Trie<Count> theTrie;
	private HashMap<Integer, HashSet<Integer>> limitedVocabs;
	private HashSet<Integer> observedVocab;
	
	public Counts() {
		theTrie = new Trie<Count>(new Count());
		observedVocab = new HashSet<Integer>();
	}
	
	public void addCount(int word, int ctx[], int count) {
		observedVocab.add(word);
		if (word == 0) {
			System.err.print("");
		}
		addCount(theTrie, word, ctx, 0, count);
	}
	
	static void addCount(Trie<Count> trie, int word, int ctx[], int idx, int count) {
		trie.getData().addCount(word, count);
		
		if (idx == ctx.length) return;
		
		// search 
		Trie<Count> subTrie = trie.getSubtrie(ctx[idx]);
		if (subTrie == null) {
			Count c = new Count();
			// add sub trie
			subTrie = new Trie<Count>(c);
			trie.addSubtrie(ctx[idx], subTrie);
		} 
		
		addCount(subTrie, word, ctx, idx+1, count);
	}
	
	public Trie<Count> getTrie() {
		return theTrie;
	}

	protected HashMap<Integer, HashSet<Integer>> getLimitedVocabs() {
		return limitedVocabs;
	}

	protected void setLimitedVocabs(HashMap<Integer, HashSet<Integer>> limitedVocabs) {
		this.limitedVocabs = limitedVocabs;
	}

	public HashSet<Integer> getObservedVocab() {
		return observedVocab;
	}

	public void setObservedVocab(HashSet<Integer> observedVocab) {
		this.observedVocab = observedVocab;
	}
}
